import os
import subprocess
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed


class Addr2LineAnalyzer:
    def __init__(self, perf_analyzer, llvm_bin_path, so_files_dir):
        self.perf_analyzer = perf_analyzer
        self.llvm_bin_path = llvm_bin_path
        self.so_files_dir = so_files_dir

    def analyze_addresses(self):
        res = []
        results = self.perf_analyzer.analyze()
        file_agg = defaultdict(list)
        for file_path, func_name, addr, event_count, percentage in results:
            so_file = file_path.split('/')[-1]
            file_agg[so_file].append((file_path, func_name, addr, event_count, percentage))

        with ThreadPoolExecutor() as executor:
            futures = [executor.submit(self._process_file, so_file, items) for so_file, items in file_agg.items()]
            for future in as_completed(futures):
                res.extend(future.result())
        return sorted(res, key=lambda item: item['event_count'], reverse=True)

    def aggregate_by_line(self, analyzed_results):
        aggregated_data = defaultdict(lambda: {'event_count': 0, 'event_percentage': 0.0})
        for result in analyzed_results:
            key = (result['source_file_path'], result['source_line'])
            aggregated_data[key]['event_count'] += result['event_count']
            aggregated_data[key]['event_percentage'] += result['event_percentage']

        aggregated_list = [{'source_file_path': k[0], 'source_line': k[1], 'event_count': v['event_count'],
                            'event_percentage': f"{v['event_percentage']:.2f}%"} for k, v in aggregated_data.items()]
        return sorted(aggregated_list, key=lambda item: item['event_count'], reverse=True)

    def aggregate_by_file(self, analyzed_results):
        file_aggregated_data = defaultdict(lambda: {'total_event_count': 0, 'total_event_percentage': 0.0, 'lines': {}})
        for result in analyzed_results:
            file_path = result['source_file_path']
            line = result['source_line']
            event_count = result['event_count']
            event_percentage = result['event_percentage']

            file_aggregated_data[file_path]['total_event_count'] += event_count
            file_aggregated_data[file_path]['total_event_percentage'] += event_percentage

            if line not in file_aggregated_data[file_path]['lines']:
                file_aggregated_data[file_path]['lines'][line] = {'event_count': 0, 'event_percentage': 0.0}

            file_aggregated_data[file_path]['lines'][line]['event_count'] += event_count
            file_aggregated_data[file_path]['lines'][line]['event_percentage'] += event_percentage

        aggregated_list = []
        for file_path, data in file_aggregated_data.items():
            lines = [{'source_line': line, 'event_count': line_data['event_count'],
                      'event_percentage': f"{line_data['event_percentage']:.2f}%"} for line, line_data in
                     data['lines'].items()]
            lines = sorted(lines, key=lambda item: item['event_count'], reverse=True)
            aggregated_list.append({
                'source_file_path': file_path,
                'total_event_count': data['total_event_count'],
                'total_event_percentage': f"{data['total_event_percentage']:.2f}%",
                'lines': lines
            })

        return sorted(aggregated_list, key=lambda item: item['total_event_count'], reverse=True)

    def _get_source_info(self, file_path, addrs):
        cmd = [os.path.join(self.llvm_bin_path, 'llvm-symbolizer.exe'), '-Cfie', file_path] + addrs
        result = subprocess.run(cmd, capture_output=True, text=True)
        return result.stdout.strip()

    def _process_file(self, so_file, items):
        addresses = [item[2] for item in items]
        source_infos = self._get_source_info(os.path.join(self.so_files_dir, so_file), addresses)
        source_infos = source_infos.split('\n\n')
        temp_res = []
        for idx in range(len(source_infos)):
            source_info = source_infos[idx]
            if '??' in source_info:
                continue
            source_func = source_info.split('\n')[0]
            source_file = ':'.join(source_info.split('\n')[-1].split(':')[:-2])
            source_line = source_info.split('\n')[-1].split(':')[-2]

            file_path, func_name, addr, event_count, percentage = items[idx]

            temp_res.append({
                'dump_file_path': file_path,
                'dump_func_name': func_name,
                'dump_addr': addr,
                'source_func': source_func,
                'source_file_path': os.path.normpath(source_file).replace('\\', '/', -1),
                'source_line': source_line,
                'event_count': event_count,
                'event_percentage': percentage,
            })
        return temp_res
